/*
 * GeoEquilibrator.cpp
 *
 *  Created on: 17 Oct 2011
 *      Author: allan
 */

#include "GeoEquilibrator.h"

// GeoReact includes
#include <Geochemistry/GeoDatabase.h>
#include <Geochemistry/GeoSystem.h>

const EquilibriumCondition
CreateSpeciesMolesConditionHelper(const string& species, const FunctionTPna& h, const Multiphase& multiphase);

const EquilibriumCondition
CreateSpeciesMolesCondition(const string& species, double moles, const Multiphase& multiphase);

const EquilibriumCondition
CreateSpeciesActivityCondition(const string& species, double activity, const Multiphase& multiphase);

const EquilibriumCondition
CreateElementMolesCondition(const string& element, double moles, const Multiphase& multiphase);

const EquilibriumCondition
CreateAcidityCondition(double pH, const Multiphase& multiphase);

const EquilibriumCondition
CreatePartialPressureGasCondition(const string& gas, double Pg, const Multiphase& multiphase);

const EquilibriumCondition
CreateChargeBalanceCondition(const Multiphase& multiphase);

GeoEquilibrator::GeoEquilibrator()
{}

GeoEquilibrator::GeoEquilibrator(
	const Multiphase& multiphase,
	const ReactionSystem& equilibrium_reactions)
: multiphase(multiphase), equilibrium_reactions(equilibrium_reactions)
{
	equilibrium_solver.Initialiaze(multiphase, equilibrium_reactions);
}

void
GeoEquilibrator::ImposeSpeciesMoles(const string& species, double moles)
{
	conditions.push_back(CreateSpeciesMolesCondition(species, moles, multiphase));
	
	composition_guess[species] = moles;
}

void
GeoEquilibrator::ImposeSpeciesActivity(const string& species, double activity)
{
	conditions.push_back(CreateSpeciesActivityCondition(species, activity, multiphase));
	
	composition_guess[species] = activity;
}

void GeoEquilibrator::ImposeElementMoles(const string& element, double moles)
{
	conditions.push_back(CreateElementMolesCondition(element, moles, multiphase));
}

void
GeoEquilibrator::ImposeAcidity(double pH)
{
	conditions.push_back(CreateAcidityCondition(pH, multiphase));
	
	composition_guess["H[+]"] = std::pow(10.0, -pH);
}

void
GeoEquilibrator::ImposePartialPressureGas(const string& gas, double partial_pressure)
{
	conditions.push_back(CreatePartialPressureGasCondition(gas, partial_pressure, multiphase));
	
	composition_guess[gas] = partial_pressure;
}

void
GeoEquilibrator::ImposeChargeBalance()
{
	conditions.push_back(CreateChargeBalanceCondition(multiphase));
}

void
GeoEquilibrator::ResetConditions()
{
	conditions.clear();
}

const ConvergenceDiagnostics
GeoEquilibrator::Equilibrate(double T, double P, VectorXd& n)
{
	EquilibriumConditions equilibrium_conditions(conditions);
	
	BOOST_FOREACH(auto pair, composition_guess)
		n[multiphase[pair.first]] = pair.second;		
	
	ConvergenceDiagnostics diagnostics = 
		equilibrium_solver.Solve(equilibrium_conditions, T, P, n);
	
	return diagnostics;
}

const ConvergenceDiagnostics
GeoEquilibrator::Equilibrate(GeoSystem& geosystem)
{
	  double T = geosystem.GetTemperature();
	  double P = geosystem.GetPressure();
	VectorXd n = geosystem.GetComposition();
	
	ConvergenceDiagnostics diagnostics = 
		Equilibrate(T, P, n);
	
	geosystem.SetState(T, P, n);
	
	return diagnostics;
}

const EquilibriumCondition
CreateSpeciesMolesConditionHelper(const string& species, 
	const FunctionTPna& h, const Multiphase& multiphase)
{
	const unsigned num_species = multiphase.GetNumSpecies();
	
	EquilibriumCondition equilibrium_condition(num_species);
	
	// Set the equilibrium condition function (h)
	equilibrium_condition.hj = h;
	
	const unsigned index_species = multiphase[species];
	
	// If the given (species) is a primary species, set the 
	// alpha_row vector of (equilibrium_condition)
	if(index_species < num_species) 
		equilibrium_condition.alpha_row[index_species] = 1.0;
	
	return equilibrium_condition;
}

const EquilibriumCondition
CreateSpeciesMolesCondition(const string& species, double moles, const Multiphase& multiphase)
{
	return CreateSpeciesMolesConditionHelper(species, ConstantTPna(moles), multiphase);
}

const EquilibriumCondition
CreateSpeciesActivityCondition(const string& species, double activity, const Multiphase& multiphase)
{
	// The index of the given species in the (multiphase) instance
	const Index i = multiphase[species];
	
	// Define the right-hand side function h = h(T, P, n, a) for this equilibrium condition
	FunctionTPna h = [=](double T, double P, const VectorXd& n, const VectorXd& a)
		{ return n[i]/a[i] * activity; };
	
	return CreateSpeciesMolesConditionHelper(species, h, multiphase);
}

const EquilibriumCondition
CreateElementMolesCondition(const string& element, double moles, const Multiphase& multiphase)
{
	// The species in the system
	const vector<string>& systemSpecies = multiphase.GetSpecies();
	 
	// The number of species in the system
	const unsigned Ni = multiphase.GetNumSpecies();
	
	// The equilibrium condition instance
	EquilibriumCondition equilibrium_condition(Ni);
	
	// Set the alphaRow entries of (equilibrium_condition) with the number of atoms of (element) of each species
	for(unsigned i = 0; i < Ni; ++i) 
		equilibrium_condition.alpha_row[i] = GeoDatabase::ElementAtoms(systemSpecies[i], element);
	
	// Set the equilibrium condition function (h)
	equilibrium_condition.hj = ConstantTPna(moles);
	
	return equilibrium_condition;
}

const EquilibriumCondition
CreateAcidityCondition(double pH, const Multiphase& multiphase)
{
	return CreateSpeciesActivityCondition("H[+]", pow(10.0, -pH), multiphase);
}

const EquilibriumCondition
CreatePartialPressureGasCondition(const string& gas, double Pg, const Multiphase& multiphase)
{
	const Index index_gaseous_phase = multiphase.FindPhase("Gaseous");
	
	const Phase& gaseous_phase = multiphase.GetPhases()[index_gaseous_phase];
	
	const vector<string>& gaseous_species = gaseous_phase.GetSpecies();
	
	const vector<Index> indexes_gaseous_species = multiphase[gaseous_species];
	
	// Define the right-hand side function h = h(T, P, n, a) for this equilibrium condition
	auto h = [=](double T, double P, const VectorXd& n, const VectorXd& a)
	{
		// The total number of moles of gaseous species
		double nG = 0.0;
		
		for(unsigned g = 0; g < gaseous_species.size(); ++g) 
			nG += n[indexes_gaseous_species[g]];
		
		return nG * Pg/P;
	};
	
	return CreateSpeciesMolesConditionHelper(gas, h, multiphase);
}

const EquilibriumCondition
CreateChargeBalanceCondition(const Multiphase& multiphase)
{
	const vector<string>& species = multiphase.GetSpecies();
	 
	const unsigned num_species = multiphase.GetNumSpecies();
	
	EquilibriumCondition equilibrium_condition(num_species);
	
	// Set the alpha_row entries of (equilibrium_condition) with the electrical charge of each species
	for(unsigned i = 0; i < num_species; ++i) 
		equilibrium_condition.alpha_row[i] = GeoDatabase::ElectricalCharge(species[i]);
	
	// Set the equilibrium condition function (h)
	equilibrium_condition.hj = ConstantTPna(0.0);
	
	return equilibrium_condition;
}

